# -*- coding: utf-8 -*-
"""
模型管理器 - 管理多个AI模型提供商
"""

import os
import json
from typing import Dict, List, Optional, Any
from ai_service_base import AIServiceBase, ModelConfig
from doubao_ai_service import DoubaoAIService
from deepseek_ai_service import DeepSeekAIService

class ModelManager:
    """AI模型管理器"""
    
    def __init__(self):
        self._services: Dict[str, AIServiceBase] = {}
        self._current_model = "doubao-seed-1-6-250615"  # 默认模型
        self._init_models()
    
    def _init_models(self):
        """初始化所有可用模型"""
        
        # 初始化豆包模型
        try:
            doubao_service = DoubaoAIService()
            self._services["doubao-seed-1-6-250615"] = doubao_service
            print("豆包模型初始化成功")
        except Exception as e:
            print(f"豆包模型初始化失败: {e}")
        
        # 初始化DeepSeek模型
        try:
            deepseek_service = DeepSeekAIService()
            self._services["deepseek-v3-1-250821"] = deepseek_service
            print("DeepSeek模型初始化成功")
        except Exception as e:
            print(f"DeepSeek模型初始化失败: {e}")
    
    def get_available_models(self) -> List[Dict[str, Any]]:
        """获取所有可用模型列表"""
        models = []
        
        for model_id, service in self._services.items():
            model_info = service.model_info.copy()
            # 检查模型是否可用（API密钥是否配置）
            model_info["available"] = self._check_model_availability(model_id)
            model_info["status"] = "available" if model_info["available"] else "need_config"
            
            # 添加模型特性信息
            if model_id == "doubao-seed-1-6-250615":
                model_info.update({
                    "description": "字节跳动豆包大语言模型，中文表现优秀",
                    "features": ["中文优化", "多轮对话", "知识问答"],
                    "icon": "🍃"
                })
            elif model_id == "deepseek-v3-1-250821":
                model_info.update({
                    "description": "DeepSeek V3推出的对话模型，代码能力强",
                    "features": ["代码生成", "逻辑推理", "数学计算"],
                    "icon": "🔍"
                })
            
            models.append(model_info)
        
        return models
    
    def _check_model_availability(self, model_id: str) -> bool:
        """检查模型是否可用"""
        if model_id not in self._services:
            return False
        
        service = self._services[model_id]
        
        # 对于豆包，检查是否有API密钥
        if model_id == "doubao-seed-1-6-250615":
            return bool(service.config.api_key)
        
        # 对于DeepSeek，检查API密钥是否配置
        elif model_id == "deepseek-v3-1-250821":
            return bool(service.config.api_key and service.config.api_key.strip())
        
        return True
    
    def set_current_model(self, model_id: str) -> bool:
        """设置当前使用的模型"""
        if model_id not in self._services:
            return False
        
        if not self._check_model_availability(model_id):
            return False
        
        self._current_model = model_id
        print(f"切换到模型: {model_id}")
        return True
    
    def get_current_model(self) -> str:
        """获取当前模型ID"""
        return self._current_model
    
    def get_current_service(self) -> Optional[AIServiceBase]:
        """获取当前模型的服务实例"""
        return self._services.get(self._current_model)
    
    async def chat_completion(self, messages: List[Dict[str, str]], model_id: str = None) -> str:
        """使用指定模型进行对话"""
        target_model = model_id or self._current_model
        
        if target_model not in self._services:
            return f"模型 {target_model} 不存在"
        
        if not self._check_model_availability(target_model):
            return f"模型 {target_model} 未配置或不可用"
        
        service = self._services[target_model]
        return await service.async_chat_completion(messages)
    
    def test_model_connection(self, model_id: str) -> Dict[str, Any]:
        """测试指定模型的连接"""
        if model_id not in self._services:
            return {
                "success": False,
                "message": f"模型 {model_id} 不存在"
            }
        
        try:
            service = self._services[model_id]
            is_connected = service.test_connection()
            
            return {
                "success": is_connected,
                "message": "连接成功" if is_connected else "连接失败",
                "model_id": model_id,
                "model_name": service.config.name
            }
        except Exception as e:
            return {
                "success": False,
                "message": f"测试连接时出错: {str(e)}",
                "model_id": model_id
            }
    
    def configure_model(self, model_id: str, api_key: str) -> Dict[str, Any]:
        """配置模型API密钥"""
        if model_id not in self._services:
            return {
                "success": False,
                "message": f"模型 {model_id} 不存在"
            }
        
        try:
            service = self._services[model_id]
            service.config.api_key = api_key.strip()
            
            # 测试配置是否有效
            test_result = self.test_model_connection(model_id)
            
            if test_result["success"]:
                print(f"模型 {model_id} 配置成功")
                return {
                    "success": True,
                    "message": f"模型 {service.config.name} 配置成功",
                    "model_id": model_id
                }
            else:
                return {
                    "success": False,
                    "message": f"配置失败: {test_result['message']}",
                    "model_id": model_id
                }
                
        except Exception as e:
            return {
                "success": False,
                "message": f"配置时出错: {str(e)}",
                "model_id": model_id
            }


# 全局模型管理器实例
model_manager = ModelManager()


def test_model_manager():
    """测试模型管理器"""
    print("=== 测试模型管理器 ===")
    
    manager = ModelManager()
    
    # 获取可用模型
    models = manager.get_available_models()
    print(f"可用模型数量: {len(models)}")
    for model in models:
        print(f"- {model['name']} ({model['id']}): {model['status']}")
    
    # 测试模型连接
    print("\n=== 测试模型连接 ===")
    for model in models:
        result = manager.test_model_connection(model['id'])
        print(f"{model['name']}: {result['message']}")


if __name__ == "__main__":
    test_model_manager()